# coding: utf-8

# @chenxin 
# @time: 2025-05-11

"""
本文件包含绘制图像压缩、视频压缩、语义通信领域模型实验结果图函数
"""

import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd 
import matplotlib
from random import choice
import random
import argparse

font = {"family": "Arial", "weight": "normal", "size": 14}
matplotlib.rc("font", **font)
LineWidth = 2

LINESTYLES = [
    "-", "--", "-.", ":"
]
MARKERS = [
    "o", "v", "^", "s", "p", "*", "h", "H", "+", "x", "D", "d",
]
COLORS = [
    [
        (0/255, 75/255, 135/255), 
        (191/255, 64/255, 64/255),
        (87/255, 135/255, 0/255) 
    ],
    [
        (0/255, 90/255, 181/255),
        (230/255, 97/255, 1/255),
        (0/255, 150/255, 130/255),
        (160/255, 0/255, 120/255)
    ],
    [
        (0/255, 114/255, 189/255),
        (217/255, 83/255, 25/255),
        (119/255, 172/255, 48/255),
        (162/255, 20/255, 47/255),
        (100/255, 100/255, 100/255),
    ],
    [
        (0/255, 92/255, 171/255),
        (106/255, 0/255, 106/255),
        (240/255, 163/255, 0/255),
        (0/255, 120/255, 120/255)
    ]
]

COLORS = sorted(COLORS, key=lambda x: len(x))


def read_data_bpp_psnr_msssim(df: pd.DataFrame, data_type: str):
    """
    data_type: psnr, msssim;
    return bpp, data, labels
    """
    num = len(df)
    bpp = []
    data = []
    labels = []
    head = "psnr" if data_type == "psnr" else "ms-ssim"
    values = df[head]
    methods = df["methods"]
    methods.fillna(False, inplace=True)
    bpps = df["bpp"]
    pre = 0
    for i in range(num):
        if methods[i]:
            labels.append(methods[i])
            if i > pre: 
                temp_val = values[pre:i].to_list() 
                temp_bpp = bpps[pre:i].to_list()
                data.append(temp_val)
                bpp.append(temp_bpp)
                pre = i
    if len(labels) != len(bpp):
        bpp.append(bpps[pre:].to_list())
        data.append(values[pre:].to_list())
    return bpp, data, labels




def plot_bpp_curve(bpp, data, labels, dataset:str, ylabel:str, save_path:str):
    """主要画 Bpp-PSNR, Bpp-MSSSIM 折线图
    bpp, data: MxN 矩阵或者列表, data 可以是 PSNR, 或者 MS-SSIM
    labels: M 个元素
    dataset: "Kadak", 在哪个数据集上做的实验
    ylabel: y 轴坐标
    save_path: 保存路径, png格式
    """
    handles = []
    num = len(labels)
    colors = [cls for cls in COLORS if len(cls) >= num][0]
    for i in range(num):
        handle, = plt.plot(bpp[i], data[i], linestyle=choice(LINESTYLES), marker=MARKERS[i], 
                           color=colors[i], linewidth=LineWidth, label=labels[i])
        handles.append(handle)

    # NOTE: 绘图
    plt.legend(handles=handles, loc=4)
    plt.grid()
    plt.xlabel("Bpp")
    plt.ylabel(ylabel=ylabel)
    plt.title(f"{dataset} datatset")
    plt.savefig(save_path, format="png", dpi=600, bbox_inches="tight")



if __name__ == "__main__":
    # NOTE: 参数化输入 
    parse = argparse.ArgumentParser()
    parse.add_argument(
        "--xlsx_path", "-xp", required=True, type=str, help="*.xlsx file path."
    )
    parse.add_argument(
        "--figure_type", "-ft", required=True, type=str, default="bpp_psnr",
        choices=["bpp_psnr", "bpp_ms-ssim"],
        help="plot figure category."
    )
    parse.add_argument(
        "--dataset", "-dt", required=True, type=str, default="Kadak",
        help="expriments on which dataset."
    )
    parse.add_argument(
        "--save_path", "-sp", default="result.png", help="the save path of png result"
    )
    parse.add_argument(
        "--seed", "-sd", default=42, help="the seed to experiment."
    )
    args = parse.parse_args()
    print(args)
    random.seed(args.seed)
    np.random.seed(args.seed)
    df = pd.read_excel(args.xlsx_path, 0)
    # NOTE: 
    if args.figure_type == "bpp_psnr":
        bpp, data, labels = read_data_bpp_psnr_msssim(df, "psnr")
        plot_bpp_curve(bpp, data, labels, args.dataset, "PSNR(dB)", args.save_path)
    elif args.figure_type == "bpp_ms-ssim":
        bpp, data, labels = read_data_bpp_psnr_msssim(df, "msssim")
        plot_bpp_curve(bpp, data, labels, args.dataset, "MS-SSIM", args.save_path)
    
    print("Successfully!!!")
    pass

